from zhipuai import ZhipuAI
import json
import mysql.connector
from datetime import datetime
import re

class ZhiPuService:

    # 发送机器人
    def send_msg(question: str):

        # res = ZhiPuService.get_history_record()
        # print(res)
        # return res + [{"role": "system", "content": 1}]

        # 
        # ===================== 第1段 =====================
        # 
        # 提问的提示词
        system_msg = ''''
            请帮我识别用户问题的意图，根据以下的关键词，判断是否需要查数据。
            关键词如下：
            -学校数
            -学生数
            -课程数
            -教师数
            -用户数
            -教育局数
        
            如果提交的问题描述与以上关键词匹配度很高时，则回复时，请遵循A的描述回答问题。
            如果提交的问题描述与以上关键词匹配度不高时，则回复时，请遵循B的描述回答问题。
            A：请回复："需要查数据"。
            B：无须按规定格式回复，请你自由回答。

        '''
        ZhiPuService.insert_record(question, 'user')
        messages=[
            {"role": "system", "content": system_msg}
        ]
        history = ZhiPuService.get_history_record() # 获取历史聊天
        messages = messages + history
        text = ZhiPuService.ai_send(messages)
        ZhiPuService.insert_record(text, 'assistant')
        print('1开始返回：' + text)
        # 判断是否要查询数据
        if text.find('需要查数据') == -1 and text.find('不需要查数据') == -1:
            return text
        
        # 
        # ==================== 第2段 ========================
        # 
        # -提示词
        system_msg = '''
            以下为表结构信息，请根据表结构信息回答问题，如果无法从表结构信息中得到答案，请说 "根据已提供的信息，无法得到答案"，不要编造答案。
            yy_school表id与yy_school_student表中的school_id关联。
            yy_school表gov_id与yy_gov表中的id关联。

            === yy_user_member 表结构 (用户表)===
            列名                   注释                             数据示例                                              
            ----------------------------------------------------------------------------------------------------
            ID                   ID                                 121                                      
            member_uid           用户唯一ID                         10245553254225418                                          
            full_name            用户姓名                           周星星                          
            phone                手机号                             13800138000                                          
            gennder              性别                               1                                      
            craeted_at           注册时间                           2023-01-01 00:00:00                                      
            ----------------------------------------------------------------------------------------------------

            === yy_school_student 表结构 (学生表)===
            列名                   注释                             数据示例                                              
            ----------------------------------------------------------------------------------------------------
            ID                   ID                                 121                                      
            gov_id               教育局ID                           10245553254225418                                          
            school_id            学校id（关联学校表id）              553              
            name                 姓名                               张三          
            sex                  性别（0未知，1男，2女）             1                                                      
            craeted_at           注册时间                           2023-01-01 00:00:00                                      
            ----------------------------------------------------------------------------------------------------

            === yy_gov 表结构 (教育局表)===
            列名                   注释                             数据示例                                              
            ----------------------------------------------------------------------------------------------------
            ID                   ID                                 121                                                                               
            name                 教育局名称                         北京市教育局                                                        
            craeted_at           注册时间                           2023-01-01 00:00:00                                      
            ----------------------------------------------------------------------------------------------------

            === yy_school 表结构 (学校表)===
            列名                   注释                             数据示例                                              
            ----------------------------------------------------------------------------------------------------
            ID                   ID                                 121                                      
            gov_id               教育局ID                           10245553254225418                                                   
            name                 学校名称                           希望小学                                                        
            craeted_at           注册时间                           2023-01-01 00:00:00                                      
            ----------------------------------------------------------------------------------------------------

            你需要根据用户提供的问题，写出对应的SQL语句。
            请注意：
            - 你只需要写出SQL语句，不需要执行。
            - 你需要根据用户提供的问题，写出对应的SQL语句，不需要解释。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的字段，则查询所有字段。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的条件，则不需要添加条件。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的排序，则不需要添加排序。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的分组，则不需要添加分组。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的限制，则不需要添加限制。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的连接，则不需要添加连接。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的子查询，则不需要添加子查询。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的聚合函数，则不需要添加聚合函数。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的去重，则不需要添加去重。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的表名，则默认查询yy_user_member表。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的字段，则查询所有字段。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的条件，则不需要添加条件。
            - 你需要根据用户提供的问题，写出对应的SQL语句，如果问题中没有提到需要查询的排序，则不需要添加排序。
            - 有关名称相关的查询，使用like模糊查询。
            请根据以上要求，写出对应的SQL语句。返回直接能执行的sql语句，不要返回解释。
        '''
        ZhiPuService.insert_record(question, 'user')
        messages=[
            {"role": "system", "content": system_msg}
        ]
        history = ZhiPuService.get_history_record() # 获取历史聊天
        messages = messages + history
        text = ZhiPuService.ai_send(messages)
        ZhiPuService.insert_record(text, 'assistant')
        print('2开始返回：' + text)
        #是否正常返回sql语句
        if text.find('select') == -1 and text.find('SELECT') == -1:
            return text

        # 清洗数据格式
        if text.find('```sql') == -1:
            sql_query = text
        else:
            pattern = r"```sql\n(.*?)\n```"
            matches = re.findall(pattern, text, re.DOTALL)
            sql_query = ''
            for match in matches:
                sql_query = match
            sql_query = sql_query.replace("```",'')
            sql_query = sql_query.replace("sql",'')
            sql_query = sql_query.replace("\n",' ')
        print('转后返回：' + sql_query)

        # 等到查询sql
        num = ZhiPuService.send_sql_data_first(sql_query)

        # 
        # ==================== 第3段 ========================
        # 
        # 提示词
        system_msg = f'''
            根据问题：{question}？此对应的SQL语句是：{sql_query};
            执行得到的结果是{num}条记录，根据此结果来回答问题。
            请注意：
            -不要使用技术语言回答问题。
            -请使用正常交谈的语气回答。
        '''
        ZhiPuService.insert_record(question, 'user')
        messages=[
            {"role": "system", "content": system_msg}
        ]
        history = ZhiPuService.get_history_record() # 获取历史聊天
        messages = messages + history
        text = ZhiPuService.ai_send(messages)
        ZhiPuService.insert_record(text, 'assistant')
        return text
    
    # 发送信息给ai
    def ai_send(messages):
        client = ZhipuAI(api_key="6bbd551b2a61b064e3d5df54448f86b9.vBPLdmoe6Pm0E2gi")  # 请填写您自己的API Key
        response = client.chat.completions.create(
            model="glm-4-0520",  # 填写需要调用的模型编码
            messages=messages,
            # stream=True,
            # response_format={ "type": "json_object" }
        )
        print(response)
        return response.choices[0].message.content

    # 增加聊天记录
    def insert_record(content, role):
        sql_query = "insert into yy_text_log (content, role, created_at) values (%s, %s, %s)"
        values = [content, role, datetime.now()]
        conn = ZhiPuService.sql_init(sql_query)
        cursor = conn.cursor()
        cursor.execute(sql_query, values)
        conn.commit()
        conn.close()

    # 获取历史聊天记录
    def get_history_record():
        sql_query = "select role,content from yy_text_log where status=1 order by id asc limit 100"
        res = ZhiPuService.send_sql_data_list(sql_query)
        messages = []
        for item in res:
            messages.append({"role": item[0], "content": item[1]})
        return messages
    
    # 查询数据表 单条
    def send_sql_data_first(sql_query: str):
        conn = ZhiPuService.sql_init(sql_query)
        cursor = conn.cursor()
        cursor.execute(sql_query)
        result  = cursor.fetchone()
        conn.close() # 关闭连接
        return result[0]
    
    # 查询数据表 多条
    def send_sql_data_list(sql_query: str):
        conn = ZhiPuService.sql_init(sql_query)
        cursor = conn.cursor()
        cursor.execute(sql_query)
        rows = cursor.fetchall()
        lists = []
        for row in rows:
            print(row)
            lists.append(row)
        conn.close() # 关闭连接
        return lists
    
    # 初始化数据库
    def sql_init(sql_query: str):
        return mysql.connector.connect(
            host='10.0.40.86',
            user='saas_upgrade',
            password='jsDwvMx8scFe3AEW',
            database='yy_saas_upgrade_test',
            port='3306'
        )

